from PIL import Image
import numpy as np
from lib.halftone.pattern import pattern
from lib.halftone.dither import dither
from lib.halftone.errdiff import error_diffusion
from lib.halftone.dotdiff import dot_diffusion
from lib.halftone.dbs import dbs

if __name__ == '__main__':
	# prepare the source image
	# image_gray = Image.open("./resources/img/painter.gif").convert("L")
	image_gray = Image.open("./resources/img/avatar.png").convert("L")

	halftone_array: list
	halftone_file: str

	print("""
	1. pattern technique
	2. dither
	3. error diffusion
	4. dot diffusion
	5. direct binary search
	""")
	n: int = int(input("input a number: "))

	if n == 1:
		halftone_array = pattern(np.array(image_gray))
		halftone_file = "./out/image_pattern.png"
	elif n == 2:
		halftone_array = dither(np.array(image_gray))
		halftone_file = "./out/image_dither.png"
	elif n == 3:
		halftone_array = error_diffusion(np.array(image_gray))
		halftone_file = "./out/image_errdiff.png"
	elif n == 4:
		halftone_array = dot_diffusion(np.array(image_gray))
		halftone_file = "./out/image_dotdiff.png"
	elif n == 5:
		halftone_array = dbs(np.array(image_gray))
		halftone_file = "./out/image_dbs.png"
	else:
		print("please input a correct number")
		assert(0)
	
	halftone_image_nparray = np.array(halftone_array, dtype=np.uint8) # make sure to add `np.uint8` to indicate a gray image
	
	print(f"original image size: { np.array(image_gray).shape }")
	print(f"image size through halftone: { halftone_image_nparray.shape }")
	
	halftone_image = Image.fromarray(halftone_image_nparray, mode="L")
	
	halftone_image.save(halftone_file)
	image_gray.save("./out/image_gray.png")

